/*
 * Copyright (c) 2018-2024 北京华恒盛世科技有限公司
 * QUICKPOOL is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/

#include <getopt.h>
#include "ee.h"

extern char **environ;

char *gstrMyName=NULL;
struct task *gptrTasks;
int gintNumTask;
JOBID_T gintExecJobId;
int gintNodeRank=-1;

char child_ee = 0;
char child_go = 0;

int debug = 0;

bool_t gboolJobMode=FALSE;
bool_t gboolTaskMode=FALSE;
int gintExecFlag=0;

struct config_param eeParams[] = {
    {"DEBUG_EE", INT_PARAM, -1, {0}},
    {"TIMEDEBUG_EE", INT_PARAM, -1, {0}},
    {NULL, STRING_PARAM, -1, {NULL}}
};

void term_handler(int);
void child_handler();
static void child_handler_ext(void);

static void init_ee_signals(void);
static void log_finish_task(int, struct task *);

static void ee_usage(char *);

int
main(int argc, char **argv)
{
    int i;

    save_daemon_path(argv[0]);

    int fordev;
    if (getenv("DEVEEDEBUG") != NULL) {

        fordev = atoi(getenv("DEVEEDEBUG"));
        fordev = MIN(120, fordev);

        for (i=0; i<fordev;i++) {
            sleep(1);
        }
    }

    char *strFromNode=NULL, *strEnvDir=NULL; 

    while (TRUE) {
        int option_index = 0;
        static struct option long_options[] = {
            {"help", no_argument, 0, 'h'},
            {"version", no_argument, 0, 'V'},
            {"dir", required_argument, 0, 'd'},
            {"node", required_argument, 0, 'm'},
            {"task", no_argument, 0, 't'},
            {"debug", no_argument, 0, '1'},
            {NULL, 0, 0, 0}
        };

        int opt;
        opt = getopt_long(argc, argv, "1Vhtfd:m:", long_options, &option_index);
        if (-1 == opt) {
            break;
        }

        switch (opt) {
        case 'd':
            strEnvDir = optarg;
            break;
        case '1':
            debug = 1;
            break;
        case 't':
            gboolTaskMode = TRUE;
            break;
        case 'm':
            strFromNode = optarg;
            gboolJobMode = TRUE;
            break; 
        case 'V':
            fputs(_QP_VERSION_, stderr);
            exit(0);
        case 'h':
	    ee_usage(argv[0]);
	    exit(0);
        case '?':
        default:
	    ee_usage(argv[0]);
	    exit(1);
        }

        if (gboolTaskMode) {
            break;
        }
    }

    char **jobExecArgv=NULL;
    if (argc > optind) {
        jobExecArgv = &argv[optind];
    }

    struct config_param logParam[] = {
        {"LOGDIR", STRING_PARAM, -1, {NULL}},
        {NULL, STRING_PARAM, -1, {NULL}},
    };

    init_config_param(logParam, strEnvDir);
    if (NULL != logParam[0].string_value) {

        char myhostnm[MAXHOSTNAMELEN+1];
        gethostname(myhostnm, MAXHOSTNAMELEN);
        if (0 == strlen(myhostnm)) {
            char *myhostname;
            if ((myhostname = qp_getmyhostname()) == NULL) {
                fprintf(stderr, "Failed to get my host name.\n");
            }
        }

        char logfile[PATH_MAX+1];
        snprintf(logfile, PATH_MAX, "%s/ee.%s.log", logParam[0].string_value, myhostnm);
        open_log(logfile, "INFO", 0);
    }

    if (init_config_param(eeParams, strEnvDir) < 0) {

        log_message(ERR, NOCLASS, "%s: init_config_param() failed due to above reason.", __func__);

        return -1;
    }

    if (gboolJobMode) {

	if (NULL == strFromNode || NULL == jobExecArgv) {
	    ee_usage(argv[0]);
	    exit(2);
	}
	 
        gintExecFlag |= JM_FLAG_STDIN | JM_FLAG_STDOUT | JM_FLAG_STDERR;
    } else if (gboolTaskMode) {
        gintExecFlag |= JM_FLAG_STDIN | JM_FLAG_STDOUT | JM_FLAG_STDERR;
    }

    gf_set_hostfile(baseParams[BASE_CONFDIR].string_value);

    close_log();

    if (debug > 1) {
    } else {

        if (NULL == gstrMyName) {

            gstrMyName = qp_getmyhostname();
            if (NULL == gstrMyName) {
                fprintf(stderr, "Failed to get my host name.\n");
                ee_quit(99);
            }
        }

        char logfile[PATH_MAX+1];
        snprintf(logfile, PATH_MAX, "%s/ee.%s.log", baseParams[BASE_LOGDIR].string_value, gstrMyName);
        int cc = open_log(logfile, baseParams[BASE_LOG_MASK].string_value, baseParams[BASE_EE_TIME_DEBUG].int_value);
        if (cc < 0) {
            fprintf(stderr, "failed to open log file[%s]: %s.\n", logfile, strerror(errno));
        }

        gf_set_logclass(baseParams[BASE_EE_DEBUG].string_value);
    }

    if (gf_chan_init() < 0) {
        return -1;
    }

    init_ee_signals();

    gintNumTask = 1;
    if (gboolJobMode || gboolTaskMode) {

        if (NULL == gstrMyName) {
            gstrMyName = qp_getmyhostname();
            if (NULL == gstrMyName) {
                log_message(ERR, NOCLASS, STR_FUNC_FAIL_M, __func__, "qp_getmyhostname");
                ee_quit(99);
            }
        }

        log_message(INFO, NOCLASS, "ee daemon started");

        JOBID_T jobId, jobIndex=0;

        char *strJobId;
        strJobId = getenv("QP_JOBID");
        if (NULL == strJobId) {
            log_message(ERR, NOCLASS, "%s: job task is not launched due to QP_JOBID not found.",
                        argv[0]);
            exit(99);
        }
        jobId = atoll(strJobId);

        char *strJobIndex;
        strJobIndex = getenv("QP_JOBINDEX");
        if (NULL != strJobIndex) {
            jobIndex = atoi(strJobIndex);
        }

        char *strNodeRank;
        strNodeRank = getenv("QP_NODE_RANK");
        if (NULL != strNodeRank) {
            gintNodeRank = atoi(strNodeRank);
        }

        if (0 < jobIndex) {
            combine_jobid(&(gintExecJobId), jobId, jobIndex);
        } else {
            gintExecJobId = jobId;
        }

        gptrTasks = (struct task *)calloc(gintNumTask, sizeof(struct task));
        if (NULL == gptrTasks) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                        "calloc", gintNumTask*sizeof(struct task *));
            ee_quit(99);
        }

        ee_executejob(jobExecArgv);
    }

    for (;;) {

        if (1 == gintNumTask && NULL != gptrTasks) {

            dump_task(&(gptrTasks[0]), 0, "one node task");

            if (0 == gptrTasks[0].end_time) {
                gf_sleep(5000);
            }
        }

        if (child_ee && child_go) {

            int numFinishTask=0;
            STATUS_WAIT_T exitStatus=0;

            for (i=0; i<gintNumTask; i++) {

                if (0 == gptrTasks[i].end_time) {
                    continue;
                }

                if (0 != WEXITSTATUS(gptrTasks[i].wait)) {
                    exitStatus = WEXITSTATUS(gptrTasks[i].wait);
                }

                numFinishTask ++;
            }

            if (numFinishTask >= gintNumTask) {

                log_message(DEBUG, EXECUTE, "%s: job ee is exiting.....", __func__);

                log_finish_task(gintNumTask, gptrTasks);

                exit(exitStatus);
            }
        }

        child_handler_ext();
    }

    return 0;
} // end function main

/**
 * @brief       quit ee.
 *
 * @param[in]   exitcode     #1: exit code
 *
 * @note        Let me log my quit reason before I quit.
 */
void
ee_quit(int exitcode)
{
    if (0 != exitcode) {
        log_message(ERR, NOCLASS, "existing with %d", exitcode);
    }

    exit(exitcode);
} // end function ee_quit

void
child_handler(void)
{
    return;
}

void
term_handler(int signum)
{
    if (gboolJobMode || gboolTaskMode) {
        return;
    }

    log_message(ERR, NOCLASS, "%s: Received signal %d, exiting", __func__, signum);

    exit(0);
}

/**
 * @brief       Handle sigchild signal.
 *
 * @note        For MPI job, gptrTasks[0] is started by ee, so we need collect
 *              its usage.
 */
void
child_handler_ext(void)
{
    int wait_flags;
    if (gboolJobMode || gboolTaskMode) {
        wait_flags = WNOHANG;
    } else {
        wait_flags = WNOHANG|WUNTRACED;
    }

    struct rusage rusage;
    memset(&rusage, 0, sizeof(struct rusage));

    int pid;
    STATUS_WAIT_T status;
    while ((pid = wait3(&status, wait_flags, &rusage)) > 0) {

        if (gptrTasks[0].pid != pid) {
            continue;
        }

        gptrTasks[0].wait = status;
        gptrTasks[0].end_time = time(NULL);

        bool_t childDone;
        if (WIFSTOPPED(status)) {
            childDone = FALSE;
        } else {
            childDone = TRUE;
        }

        if (childDone) {

            double ru_utime;
            ru_utime = rusage.ru_utime.tv_sec + rusage.ru_utime.tv_usec / 1000000.0;

            double ru_stime;
            ru_stime = rusage.ru_stime.tv_sec + rusage.ru_stime.tv_usec / 1000000.0;

            gptrTasks[0].inst_usage->utime = ru_utime;
            gptrTasks[0].inst_usage->stime = ru_stime;
        }
    }

    return;
} // end function child_handler_ext

/* ------------------------------------------------
 *
 *  below are static functions used in this file
 *
 * ------------------------------------------------ */

static void
init_ee_signals(void)
{
    Signal_(SIGCHLD, (SIGFUNCTYPE) child_handler);
    Signal_(SIGINT, (SIGFUNCTYPE) term_handler);
    Signal_(SIGHUP, SIG_IGN);
    Signal_(SIGPIPE, SIG_IGN);
    Signal_(SIGTTIN, SIG_IGN);
    Signal_(SIGTTOU, SIG_IGN);
    Signal_(SIGTSTP, SIG_IGN);
    Signal_(SIGCONT, SIG_IGN);

#ifdef SIGDANGER
    Signal_(SIGDANGER, SIG_IGN);
#endif

    Signal_(SIGTERM, (SIGFUNCTYPE) term_handler);
#ifdef SIGXCPU
    Signal_(SIGXCPU, (SIGFUNCTYPE) term_handler);
#endif

#ifdef SIGXFSZ
    Signal_(SIGXFSZ, (SIGFUNCTYPE) term_handler);
#endif

#ifdef SIGPROF
    Signal_(SIGPROF, (SIGFUNCTYPE) term_handler);
#endif

#ifdef SIGLOST
    Signal_(SIGLOST, (SIGFUNCTYPE) term_handler);
#endif

    Signal_(SIGUSR1, (SIGFUNCTYPE) term_handler);
    Signal_(SIGUSR2, (SIGFUNCTYPE) term_handler);
#ifdef SIGABRT
    Signal_(SIGABRT, (SIGFUNCTYPE) term_handler);
#endif

    Signal_(SIGQUIT, SIG_IGN);

    return;
} // end function init_ee_signals

/**
 * @brief       Write job task usage into job task acct file.
 *
 * @param[in]   numtask      #1: number of total task of the job
 * @param[in]   tasks        #2: tasks of the job
 *
 * @note        Only log acct for the job who runs on multiple nodes.
 *              JM on first node would read this file to collect job usage
 *              on different nodes.
 */
void
log_finish_task(int numtask, struct task *tasks)
{
    char *tmpDir = getenv("TMPDIR");
    if (NULL == tmpDir) {
        log_message(ERR, NOCLASS, "%s: do not where to log finish task.", __func__);
        return;
    }

    char taskAcctFile[PATH_MAX+1];
    snprintf(taskAcctFile, PATH_MAX, "%s/%lld.ee.acct", tmpDir, gintExecJobId);

    FILE *fileAcct;
    fileAcct = fopen(taskAcctFile, "w");
    if (NULL == fileAcct) {
        log_message(ERR, NOCLASS, STR_FUNC_S_FAIL_M, __func__, "fopen", taskAcctFile);
        return;
    }

    struct record taskRecord;
    snprintf(taskRecord.version, MAXVERSIONLEN, "%d", QP_XDR_VERSION);
    taskRecord.type = RECORD_TYPE_JOB_ACCT;
    taskRecord.log_time = time(NULL);

    unsigned int num = 2; // gintExecJobId, numtask

    num += 7*numtask; // node_rank, pid, pgid, wait, exit_code, start_time, end_time

    int i;
    for (i=0; i<numtask; i++) {

        if (NULL == tasks[i].inst_usage) {
            continue;
        }

        num += 6; //mem, swap, utime, stime, npids, npgids

        num += tasks[i].inst_usage->npids; //pid, ppid, pgid
        num += tasks[i].inst_usage->npgids; //pgid
    }

    taskRecord.KVs = (key_value_t *)calloc(num, sizeof(key_value_t));
    if (NULL == taskRecord.KVs) {
        log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                    "calloc", num*sizeof(key_value_t));
        return;
    }

    num = 0;
    taskRecord.KVs[num].key = RECORD_KEY_JOBID;
    taskRecord.KVs[num].value = (char *)calloc(32, sizeof(char));
    if (NULL == taskRecord.KVs[num].value) {
        log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 32);
        return;
    }
    snprintf(taskRecord.KVs[0].value, 32, "%lld", gintExecJobId);
    num ++;

    taskRecord.KVs[num].key = RECORD_KEY_NUM_TASK;
    taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
    if (NULL == taskRecord.KVs[num].value) {
        log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
        return;
    }
    snprintf(taskRecord.KVs[num].value, 12, "%d", numtask);
    num ++;

    for (i=0; i<numtask; i++) {

        taskRecord.KVs[num].key = RECORD_KEY_NODE_RANK;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", gintNodeRank);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_JOB_PID;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].pid);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_JOB_PGID;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].pgid);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_TASK_STATUS;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].wait);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_EXIT_CODE;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].exit_code);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_START_TIME;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].start_time);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_END_TIME;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].end_time);
        num ++;

        if (NULL == tasks[i].inst_usage) {
            continue;
        }

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_MEM;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].inst_usage->mem);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_SWAP;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].inst_usage->swap);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_UTIME;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].inst_usage->utime);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_STIME;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%ld", tasks[i].inst_usage->stime);
        num ++;

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_NPID;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].inst_usage->npids);
        num ++;

        unsigned int j;
        for (j=0; j<tasks[i].inst_usage->npids; j++) {
            taskRecord.KVs[num].key = RECORD_KEY_USAGE_PIDINFO;
            taskRecord.KVs[num].value = (char *)calloc(36, sizeof(char));
            if (NULL == taskRecord.KVs[num].value) {
                log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 36);
                return;
            }

            snprintf(taskRecord.KVs[num].value, 36, "%d %d %d",
                     tasks[i].inst_usage->pidinfo[j].pid,
                     tasks[i].inst_usage->pidinfo[j].ppid,
                     tasks[i].inst_usage->pidinfo[j].pgid);
            num ++;
        }

        taskRecord.KVs[num].key = RECORD_KEY_USAGE_NPGID;
        taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
        if (NULL == taskRecord.KVs[num].value) {
            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
            return;
        }
        snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].inst_usage->npgids);
        num ++;

        for (j=0; j<tasks[i].inst_usage->npgids; j++) {
            taskRecord.KVs[num].key = RECORD_KEY_USAGE_PGID;
            taskRecord.KVs[num].value = (char *)calloc(12, sizeof(char));
            if (NULL == taskRecord.KVs[num].value) {
                log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__, "calloc", 12);
                return;
            }
            snprintf(taskRecord.KVs[num].value, 12, "%d", tasks[i].inst_usage->pgid[i]);
            num ++;
        }
    }

    taskRecord.numKVs = num;

    if (write_record(fileAcct, &taskRecord) == -1) {
        fclose(fileAcct);
        return;
    }

    fclose(fileAcct);

    return;
} // end function log_job_acct

static void
ee_usage(char *cmd)
{
    fprintf(stdout, ": %s  [-V] [-h] [-debug_level] [-d env_dir] [-m cl_host command [args]]\n",
            cmd);

    return;
}
